package smat.test

import smat.intr.Matrix
import smat.impl.MatrixJC

class Discrete_Bayes {
  val all = Matrix.all
  
  def discreteBayes(trainPatterns: Matrix, trainTargets: Matrix, testPatterns: Matrix): Matrix = {
    val uc = trainTargets.unique
    var n = 1
    var uniquePatterns = new MatrixJC(0, 0)
    val counts = new MatrixJC(0, 0)
    for(i <- trainPatterns.cols){
      val data = trainPatterns(all, i)
      val patMat = new MatrixJC(trainPatterns.size:_*)
      trainPatterns.cols.foreach(col => patMat((all, col)) = data)
      val indices = (patMat - trainPatterns).sum.find(0)
      if(uniquePatterns.isEmpty){
        uniquePatterns((data.indices, 0)) = data
        uc.indices.foreach(j => counts((j, 0)) = trainTargets(indices).count(uc(j)))
      } else{
        val uqMat = new MatrixJC(uniquePatterns.size:_*)
        uniquePatterns.cols.foreach(col => uqMat((all, col)) = data)
        if((uqMat - uniquePatterns).count(0) == 0){
          uniquePatterns((all, n)) = data
          uc.indices.foreach(j => counts((j, n)) = trainTargets(indices).count(uc(j)))
          n += 1
        }
      }
    }
    val sc = counts.sum
    val mat = new MatrixJC(counts.size:_*)
    uc.indices.foreach(line => mat((line, all)) = sc)
    val pX = counts / mat
    val uniqueTargets = MatrixJC(pX.cols.map(col => pX((all, col)).maxIndex))
    
    return nearestNeighbor(uniquePatterns, uniqueTargets, testPatterns, 1)
  }
  
  def nearestNeighbor(trainPatterns: Matrix, trainTargets: Matrix, testPatterns: Matrix, knn: Int): Matrix = {
    val uc = trainTargets.unique
    
    if(trainTargets.size(0) < knn){
      throw new IllegalArgumentException("You specified more neighbors than there are points.")
    }
    
    val n = testPatterns.size(1)
    val testTargets = MatrixJC.zeros(n)
    for(i <- testPatterns.cols){
      val mat = new MatrixJC(trainTargets.size(0), n)
      val vect = testPatterns(all, i)
      trainTargets.cols.foreach(col => mat((all, col)) = vect)
      val diff = trainPatterns - mat
      val dist = (diff * diff).sum
      val indices = dist.unique.flatMap(e => dist.find(e))
      val filter = trainTargets(indices.take(knn))
      val hist = uc.map(e => filter.count(e))
      val best = MatrixJC(hist).maxIndex
      testTargets(i) = uc(best)
    }
    return testTargets
  }
}
